import logging
from typing import Callable, Union, List
import dspy
import requests
import re

import os
import logging
from typing import Union, List, Callable
import requests

# 假设您有一个 dspy.Retrieve 父类，如没有可自定义
import dspy


class PaperRM(dspy.Retrieve):
    def __init__(
        self,
        k: int = 3,
        default_top_k: int = 30,
        is_valid_source: Callable[[str], bool] = None
    ):
        """
        初始化检索模块，用于通过论文搜索接口来检索论文片段。
        
        Args:
            k (int): 在 forward 中决定返回的前 k 条结果（若需要限制）。默认为 3.
            default_top_k (int): 在请求接口时的 top_k 参数默认值，默认为 30.
            is_valid_source (Callable[[str], bool], optional): 
                用于对返回结果进行过滤的函数。如不需要过滤，可不传或传 None.
        """
        # 让父类先初始化, 内部可能会记录 self.k
        super().__init__(k=k)
        
        # 在请求接口时, 若调用者不额外指定 top_k, 就会用这个默认值
        self.default_top_k = default_top_k

        # 如果你还想统计调用次数
        self.usage = 0

        # 如果给了过滤函数就用，否则默认不过滤
        if is_valid_source:
            self.is_valid_source = is_valid_source
        else:
            self.is_valid_source = lambda x: True

        logging.info(
            f"PaperRM initialized with k={k}, default_top_k={default_top_k}, "
            f"is_valid_source={is_valid_source is not None}"
        )

    def get_usage_and_reset(self):
        """
        读取 usage 并清零，可用于统计一段时间内的调用次数。
        """
        usage = self.usage
        self.usage = 0
        return {"PaperRM": usage}

    def forward(
        self,
        query_or_queries: Union[str, List[str]],
        top_k: int = None,



        ####此处参数未做处理
        ####后续需要改进
        exclude_urls: List[str] = None  # 新增参数
    ):
        """
        调用论文搜索接口，根据文本查询搜索论文片段。
        
        Args:
            query_or_queries (Union[str, List[str]]): 查询或查询列表
            top_k (int, optional): 请求接口时返回的条数，不传则用默认值 self.default_top_k
        
        Returns:
            list: 包含检索结果的列表，每个元素是一个 dict，表示一条论文片段信息。
        """
        # 如果只传入一个字符串，就封装成列表
        queries = [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries
        
        # 统计 usage
        self.usage += len(queries)
        
        # 如果调用者没传 top_k，就用默认值
        if top_k is None:
            top_k = self.default_top_k
        
        collected_results = []
        for query in queries:
            try:
                # 组装请求参数
                params = {
                    "query": query,
                    "top_k": top_k
                }
                # 发起 HTTP GET 请求
                resp = requests.get(
                    "http://180.184.65.98:38880/atomgit/search_papers",
                    params=params,
                    timeout=15
                )
                # 如果状态码非200，抛出异常
                resp.raise_for_status()

                # 根据题目说明，这里返回一个 JSON 数组
                raw_results = resp.json()  # 返回形如 [{"id":..., "distance":..., "entity":{...}}, ...]

                # change add reranker
                from .storm_wiki.modules.rerank import rm_rerank

                new_order_results = rm_rerank(query, raw_results)
                raw_results = new_order_results

                for raw_item in raw_results:
                    entity = raw_item.get("entity", {})
                    #在此处对于entity进行拼接，获取完整的description引用
                    original_filename = entity.get("original_filename")
                    match = re.search(r"^(.*?)_Paper_Meta_Data_(.*?)_with_whole_text\.db", original_filename)
                    year = entity.get("year")
                    if match:
                        type_ = match.group(1)
                        name_ = match.group(2)
                        if type_ == "Conf":
                            name_ = re.sub(r'\d+', '', name_.split("_")[0])  
                        elif type_ == "Journal":
                            name_ = name_.replace("_", " ")
                        description = f"{name_},{year},chunk {entity.get('chunk_id')}"  # 格式化为期望的格式
                    else:
                        if year:
                            description = f",未知会议/期刊,{year},chunk {entity.get('chunk_id')}"
                        else:
                            description = f",未知会议/期刊,未知年份,chunk {entity.get('chunk_id')}"
                    
                    # 构造与 YouRM 类似的 dict
                    converted = {
                        "title": entity.get("paper_title"),
                        "url": str(entity.get("paper_id")) + "@chunkID" + str(entity.get("chunk_id")),

                        #对于格式进行二次调整：description中加入期刊以及年份进行拼接
                        #最终具体引用格式如下：[xx]title，期刊名，年份，chunk x
                        "description": description,
                        # snippets我们可以让它是一个列表，包含 chunk_text
                        "snippets": [entity.get("chunk_text")] if entity.get("chunk_text") else []
                    }
                    collected_results.append(converted)
                
            except Exception as e:
                logging.error(f"Error occurs when searching query={query}: {e}")
        
        return collected_results